package com.yifeng.repo.controller.rocketmq.message.produce;

import lombok.AllArgsConstructor;
import lombok.Data;
import org.apache.rocketmq.client.producer.MessageQueueSelector;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageQueue;

import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Created by daibing on 2023/8/5.
 */
public class TargetMessageQueueSelector implements MessageQueueSelector {
    private final ConcurrentMap<Object, ArgData> arg2Data = new ConcurrentHashMap<>();

    @Override
    public MessageQueue select(List<MessageQueue> mqs, Message msg, Object arg) {
        // 对靶向参数arg做负载均衡
        ArgData argData = getArgData(arg);
        int index = Math.abs(argData.getLoadBalanceIndex().getAndAdd(1) % argData.getQueueIndexes().length);

        // 根据分配到的队列序号获取实际队列
        return mqs.get(Integer.parseInt(argData.getQueueIndexes()[index]));
    }

    private ArgData getArgData(Object arg) {
        ArgData argData = arg2Data.get(arg);
        if (argData != null) {
            return argData;
        }
        synchronized (arg2Data) {
            argData = arg2Data.get(arg);
            if (argData != null) {
                return argData;
            }
            argData = new ArgData(String.valueOf(arg).split(","), new AtomicInteger(0));
            arg2Data.put(arg, argData);
            return argData;
        }
    }

    @Data
    @AllArgsConstructor
    private static class ArgData {
        // 分割arg数字拼接字符串，得到多个队列索引序号
        private String[] queueIndexes;
        // 负载均衡计数器
        private AtomicInteger loadBalanceIndex;
    }
}
